import threading
import datetime
import uuid
from typing import List, Dict, Any, Optional
from database.milvus_client import MilvusClient
from pymilvus import FieldSchema, DataType
import logging
from models.memory_record import MemoryRecord, ContentType, Role
from utils.aes_util import AES256Cipher

logger = logging.getLogger(__name__)


class MemoryRecordServer(MilvusClient):
    _instance = None
    _lock = threading.Lock()

    # 类常量定义
    COLLECTION_NAME = "memory_records_s_4"
    COLLECTION_FIELDS = [
        FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=36, is_primary=True, auto_id=False),
        FieldSchema(name="user_id", dtype=DataType.VARCHAR, max_length=36, default_value=""),
        FieldSchema(name="saas_id", dtype=DataType.VARCHAR, max_length=36, default_value=""),
        FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535, default_value=""),
        FieldSchema(name="type", dtype=DataType.VARCHAR, max_length=255, default_value=""),
        FieldSchema(name="sec_type", dtype=DataType.VARCHAR, max_length=255, default_value=""),
        FieldSchema(name="optimized_content", dtype=DataType.VARCHAR, max_length=65535, default_value=""),
        FieldSchema(name="relation", dtype=DataType.VARCHAR, max_length=65535, default_value=""),
        FieldSchema(name="timestamp", dtype=DataType.VARCHAR, max_length=64, default_value=""),
        FieldSchema(name="happen_time", dtype=DataType.VARCHAR, max_length=64, default_value=""),
        FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=255, default_value=""),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=2560),
        FieldSchema(name="detail", dtype=DataType.VARCHAR, max_length=65535, default_value=""),
    ]
    DEFAULT_SEARCH_PARAMS = {
        "metric_type": "L2",
        "params": {"nprobe": 16}
    }

    def __new__(cls, *args, **kwargs):
        """线程安全的单例模式实现"""
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance

    def __init__(self, host="localhost", port="19530", user="root", pwd="szsk8888"):
        """初始化方法，确保只执行一次"""
        if self._initialized:
            return

        # AES加密初始化
        password = "szsk"
        salt = b'fixed_salt_12345'
        self.cipher = AES256Cipher(password, salt=salt)

        # 父类初始化
        super().__init__(
            host=host,
            port=port,
            collection_name=self.COLLECTION_NAME,
            fields=self.COLLECTION_FIELDS,
            user=user,
            pwd=pwd
        )

        self._initialized = True

    def _prepare_record_data(self, record: MemoryRecord, embedding: list) -> list:
        """准备要插入的记录数据"""
        record.timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        return [
            [str(uuid.uuid4())],
            [record.user_id],
            [record.saas_id],
            [self.cipher.encrypt(record.content)],
            [record.type],
            [record.sec_type],
            [self.cipher.encrypt(record.optimized_content)],
            [record.relation],
            [record.timestamp],
            [record.happen_time],
            [record.title],
            [embedding],
            [record.detail]
        ]

    def insert_record(self, record: MemoryRecord, embedding: list) -> str:
        """插入新记录"""
        try:
            start_time = datetime.datetime.now()
            logger.debug(f"开始插入记录: {record.id}")

            data = self._prepare_record_data(record, embedding)
            self.collection.insert(data)

            logger.debug(
                f"成功插入记录: {record.id} [耗时: {(datetime.datetime.now() - start_time).total_seconds():.2f}s]")
            return record.id
        except Exception as e:
            logger.error(f"插入记录失败: {str(e)}", exc_info=True)
            raise RuntimeError(f"插入记录失败: {str(e)}") from e

    def _decrypt_record(self, hit_entity: Dict[str, Any]) -> Dict[str, Any]:
        """解密记录内容"""
        return {
            "relation": hit_entity.get('relation'),
            "content": self.cipher.decrypt(hit_entity.get('content')),
            "optimized_content": self.cipher.decrypt(hit_entity.get('optimized_content')),
            "title": hit_entity.get('title'),
            "type": hit_entity.get('type'),
            "sec_type": hit_entity.get('sec_type'),
            "detail": hit_entity.get('detail')
        }

    def search_records(
            self,
            user_id: str,
            saas_id: str,
            query_embedding: list,
            content_type: Optional[ContentType] = None,
            role: Optional[Role] = None,
            top_k: int = 1,
            min_score: float = 1.0,
            date: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """搜索记录"""


        try:
            if not date:
            # 构建查询表达式
                expr = f'user_id == "{user_id}" and saas_id == "{saas_id}"'

                results = self.collection.search(
                    data=[query_embedding],
                    anns_field="embedding",
                    param=self.DEFAULT_SEARCH_PARAMS,
                    limit=top_k,
                    expr=expr,
                    output_fields=["relation", "content", "optimized_content", "title", "type", "sec_type","detail"]
                )
                records = []
                for hits in results:
                    for hit in hits:
                        if hit.score > min_score:
                            continue
                        logger.debug(f"检索到记录，相似度分数: {hit.score:.4f}")
                        records.append(self._decrypt_record(hit.entity))
            else:
                expr = f'user_id == "{user_id}" and happen_time == "{date}"'
                #expr = f'user_id == "{user_id}" and saas_id == "{saas_id}"'

                results = self.collection.search(
                    data=[query_embedding],
                    anns_field="embedding",
                    param=self.DEFAULT_SEARCH_PARAMS,
                    limit=top_k,
                    expr=expr,
                    output_fields=["relation", "content", "optimized_content", "title", "type", "sec_type", "detail"]
                )
                records = []
                for hits in results:
                    for hit in hits:


                        records.append(self._decrypt_record(hit.entity))

            return records
        except Exception as e:
            logger.error(f"搜索记录失败: {str(e)}", exc_info=True)
            raise RuntimeError(f"搜索记录失败: {str(e)}") from e